"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$$HeadURL$$

Created on:           7/29/2009
Last Revised on:      $$Date$$
Version:              $$Revision$$
ID:                   $$Id$$
Author:               $$Author$$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================
Updated 6/25/2009: added support for multi-resolution (coarse to fine) registration.
Updated 8/10/2009: added support for multi-iteration registration
"""

#===============================================================================
# Imports
#===============================================================================

import JT

import numpy as np
import scipy.ndimage as ndi
import pylab
import time
import os
import csv

time_last = time.time()

class Registration():
    """
    Usage:
    either call: SetFixedFileName(pathToImage) or SetFixedImage(vtkImage)
    then call either: SetMovingFileName(pathToImage) or SetMovingImage(vtkImage)
    a call to GetValue(), will return the metric value.
    """

    def __init__(self):
        """
        """
        self._firstRun = True
        self._fixedImage = None
        self._movingImage = None

        #Specify the region of interest
        self._roi = None
        
        #A matplotlib figure to display fixed and moving images.
        self._imageDisplayPanel = None
        self._IC = JT.ImageCanvas()
        
        #Iteration counter
        self._iterCount = 0
        
        self._plot = False
        self._debug = True
        self._logGradientResult = True
        
        # Setup default image metric
        self._imageMetric = JT.Metric
        # Define a dictionary to store metric parameters.
        # The exact parameters required will vary depending on the chosen metric
        self._metricParams = dict()
        self._metricParams['method'] = 'bpc'
        
        # Currently optimizer parameters are not passed using this dictionary!
        # But in the future this is going to be the preferred method.
        self._optimizerParams = dict()
        
        self._val_log = np.array([])
        self._pose_log = np.ones(6)
        self._data = []
        self.bestCost = 1.0e300
        self.bestPose = None
        
        # Setup a default optimizer
        #self.optimizer = JT.Optimizer.RegularStepGradientDescentOptimizer()
        # Optimizer Parameters:
        self.optimizer = None
        self._initialPose = None
        self._initialStartingPose = None
        self._scales = None
        
        self._numberOfParameters = 6
        
        #self.StartNewLog()
        
        
    def StartNewLog(self):
        """
        # Setup a log file for saving the pose and metric value.
        # TODO: Eventually may want to extract this to a place where it can 
        # be reused
        """
        logTag = 'JT_Results'     #This tag can be used to identify trials
        logName = time.strftime('%Y%m%d_%H%M%S')+'_'+logTag+'.log'
        logDir = os.path.join(os.getcwd(),'log')
        self.logFile = open(os.path.join(logDir,logName),'wt')
        print 'Logging dir: ', logDir  
        try:
            os.mkdir(logDir)
        except:
            pass
        self.log = csv.writer(self.logFile,dialect='excel-tab')
        # Write Headers
        self.log.writerow(('Iteration','Cost','Pose0(mm)','Pose1(mm)','Pose2(mm)',
                                'Pose3(deg)','Pose4(deg)','Pose5(deg)','BestCost'))
    
    def GetNumberOfParameters(self):
        """For ITK Optimizers"""
        return self._numberOfParameters

    def SetMovingImageGenerator(self,ImageGenerator):
        """
        Set an object that can generate a 2D moving image.
        Currently, this is limited to the DRR class.
        IMPORTANT: fixed image must be set first so that the DRR image size
        can be properly set.
        """
        if self._fixedImage == None:
            print "Error: Fixed image must be set first"
            return
        
        self._imageGenerator = ImageGenerator
        self._imageGenerator.SetRenderWindowSize(self._originalFixedImageSize)


    def UpdateImageGenerator(self):
        """
        Call the GenerateImage function and save the resulting image as the 
        new moving image.
        """
        self._movingImage = self._imageGenerator.GetImage()

        #Set the moving image spacing to be the same as fixed image spacing
        #Is this really necessary each time an image is generated???
        if self._fixedImage != None:
            self._movingImage.SetSpacing(self._fixedImage.GetSpacing())
            
    
    def SetFixedImage(self,fixedImage):
        """
        Set a VTK image.  The expected usage is:
        reg.SetFixedImage(FixedImage.GetImage(i))
        
        It will be converted to a numpy image in the GetValue method.
        """
        self._fixedImage = fixedImage
        self._originalFixedImageSize = fixedImage.GetDimensions()[0:2]
        
        

    def SetImageMetric(self,metric):
        """
        CHANGED: specify metric as string {'nmi', 'ncc', 'mi', 'necc', 'ecc', 'mse', 'gd'}
        This is passed to the Metric module in the parameter dictionary.
        """
        #self._imageMetric = metric
        self._metricParams['method']=metric
    
    def SetOptimizer(self,optimizer):
        """
        Connect the optimizer to the registration class.  This interface is 
        in flux, beware of changes.
        
        The current implementation is for the user to setup all registration 
        components (metric, optimizer, images) in the registration class 
        and then call 'Start' (in the registration class) to initiate the 
        registration.  This will in turn call the optimize function.
        """
        self.optimizer = optimizer
    
    def GetOptimizer(self):
        """
        Return the current optimizer
        """
        return self.optimizer
    
    
    def SetInitialPose(self,pose):
        """
        Pose is a 6 element vector (Tx, Ty, Tz, Rx, Ry, Rz).  The initial 
        pose is used to give the optimizer a strarting position.  It can be 
        determined from a specified value (for benchmarking), or through 
        user interaction (KneeTrack, JointTrack).
        """
        self._initalPose = pose
        try:
            self.optimizer.SetInitialPosition(pose)
        except AttributeError:
            print "Error! Optimizer must be set before pose can be initialized."
    
    def SetScales(self,scales):
        """
        Scales is a 6 element vector.  Scales is used to weight the gradient
        so that the rotation scales are approximately equivalent to 
        translation scales.
        """
        self._scales = scales
        self.optimizer.SetScales(scales)
        
    
    def SetRegionOfInterest(self,roi):
        """
        Set a region of interest for to apply the metric to.  This could be 
        the same region specified in a image mask.  This roi will be transformed
        (scaled) by the resolution factor so that it will match the image at
        the desired resolution.
        """
        self._roi = roi
        self._transformedRoi = roi
        self._metricParams['roi'] = roi
    
    def Start(self):
        """
        Start the registration process by kicking off the optimzer
        """
        self.StartNewLog()
        self.Reset()
        self.optimizer.StartOptimization()
        self.OptimizerDone()
    
    def OptimizerDone(self):
        """
        Perform any clean up functions necessary after the optimizer is 
        finished running.
        """
        # Make sure that the registration did not find a better solution 
        # while doing the gradient descent calculations
        if self.bestCost < self.optimizer._Value:
            print "The optimizer returned a value of: ", self.optimizer._Value
            print "But this was not the best solution found."
            print "Which was: ", self.bestCost
            print "At Pose  : ", self.bestPose
        else:
            self.bestCost = self.optimizer._Value
            self.bestPose = self.optimizer._CurrentPosition
        
        self.logFile.close()
        try:
            pylab.close(self._imageDisplayPanel)
            self._imageDisplayPanel = None
        except:
            print "No display exists. Continuing happily"
    
    def SmoothImage(self,image):
        """
        Apply a VTK median filter to the given VTK image and return the resultant image.
        the size of the kernel is set using self._factor and is implemented here 
        as the same for both X and Y 2D images.
        """
        print "Images smoothed, radius: ", self._factor
        smoother = JT.vtk.vtkImageMedian3D()
        smoother.SetInput(image)
        smoother.SetKernelSize(self._factor, self._factor, 1)
        smoother.Update()
        
        # Alternatice implementation
#        if self._useSmoothImages:
#            image1 = self.SmoothImage(self._fixedImage)
#            image2 = self.SmoothImage(self._movingImage)
#        else:
#            image1 = self._fixedImage
#            image2 = self._movingImage
            
        return smoother.GetOutput()

    def GetValue(self,pose,volume=0,log_result=True):
        """
        return the current value of the image metric.  Set the position and
        orientation of the transform.  
        INPUTS:
            Pose - vector made up of (position[3],orientation[3])
            volume - the (integer) index of item_list, to which the pose will be applied 
                     allows multiple volumes to be registered to the same fixed image.
        OUTPUT:
            value - scalar cost function evaluation.  Cost function can be set 
                    using my Metric.py set of functions or to any other function
                    that can take two images and return a scalar cost. 
        """
        global time_last

        #Check to make sure a 6 element pose was supplied
        try:
            assert len(pose)==6
        except AssertionError:
            print "Expected 'pose' to contain 6 elements, got %i" % len(pose)
            raise
    
        # First update the pose of the volume before the image is generated
        # WARNING: This interface might be changed to better manage multiple volumes
        self._imageGenerator._volume.SetPose(pose)
        self.UpdateImageGenerator()
        
        image1 = self._fixedImage
        image2 = self._movingImage
        
        # Check components to make sure everything is up-to-date
        # Returns false if there is a problem with the components
        if self.Update() == False:
            print "Could not update components, nothing done"
            raise AssertionError
        
        # Convert VTK images to numpy arrays for image metric calculation
        image1 = JT.utils.vtk2numpy(image1)
        image2 = JT.utils.vtk2numpy(image2)
        
        # This is the key metric evaluation step
        try:
            val=self._imageMetric.GetValue(image1,image2,**self._metricParams)
            self._iterCount += 1
        except:
            print "Error calculating metric"
            raise
                
        if val < self.bestCost:
            self.bestCost = val
            self.bestPose = pose

        if self._debug and log_result:
            # only write results if debug is set and log_result is true
            # log_result is turned off when gradients are calculated
            print "%i(%fs), Metric:,%4.4f, Pose[mm,deg]:, %4.2f, %4.2f, %4.2f, %4.4f, %4.4f, %4.4f"%\
                   (self._iterCount,time.time()-time_last,val,pose[0],pose[1],
                    pose[2],np.degrees(pose[3]),np.degrees(pose[4]),np.degrees(pose[5]))
            time_last = time.time()
            # Log status/results to a file that can be plotted later.
            data_list = [self._iterCount, val, pose[0], pose[1], pose[2],
                        np.degrees(pose[3]), np.degrees(pose[4]),
                        np.degrees(pose[5]), self.bestCost]
            self._data.append(data_list)
            self.log.writerow(data_list)
        
        if self._plot and log_result:
            self.ShowImages(image1,image2)
        

        return val
    
    def GetGradient(self,pose,step_length=None):
        """
        Inputs
        pose is a 6 element vector describing the 3 translation and 3 rotations
        of the 3D object.
        
        step is a scalar float describing the desired step size.  It is 
        expected that translations and rotations could have very different 
        step sizes (mm vs radians).  Currently, step size is applied in terms
        of mm for translations and degrees for rotations
        
        Returns
        gradient is a 6 element gradient vector that represents the difference 
        between a positive and negative step.
        """
        cost1 = np.zeros(len(pose))
        cost2 = np.zeros(len(pose))
        
        if step_length == None:
            step = np.array([1,1,1,np.pi/90,np.pi/90,np.pi/90])
        else:
            step = step_length * np.array([1,1,1,np.pi/180,np.pi/180,np.pi/180])
        
        for i,p in enumerate(pose):
            # Find the cost at negative step
            x1 = pose.copy()
            x1[i] = x1[i] - step[i]
            cost1[i] = self.GetValue(x1,log_result=self._logGradientResult)
            
            # Find the cost at positive step
            x2 = pose.copy()
            x2[i] = x2[i] + step[i]
            cost2[i] = self.GetValue(x2,log_result=self._logGradientResult)
            
        gradient = (cost2 - cost1) / (2 * step)
        
        #print "Gradient: ", gradient
        return gradient
        
    
    def ShowImages(self,image1,image2):
        """
        Create a matplotlib figure and display the current fixed image and
        moving image.  Images will be cropped to the specified ROI.
        """
        # Only display the region of interest from the image.
        roi = self._transformedRoi
        pylab.ion()
        pylab.bone()

        # The first time this function is run, create the figure and add the fixed image
        if self._imageDisplayPanel == None:
            print "Creating Display Panel"
            self._imageDisplayPanel = pylab.figure(figsize=(10,2.5),
                                                   facecolor=(1,1,1))
            self._imageDisplayPanel.subplots_adjust(bottom=0.05,top=0.95,
                                                    hspace=0.1,wspace=0.1)
            # Add fixed image (may not be needed if fixed image is not changing 
            # after creation of display panel
            sp1 = self._imageDisplayPanel.add_subplot(121,xticks=[],yticks=[],
                                                      frameon=False)
            #sp1.cla() 
            image1 = image1[roi[0][0]:roi[1][0], roi[0][1]:roi[1][1]]
            sp1.imshow(image1,origin='lower')
            #print "Adding image1"
        
        # Add moving image, definately will need to be updated with each iteration
        sp2 = self._imageDisplayPanel.add_subplot(122,xticks=[],yticks=[],
                                                  frameon=False)
        #sp2.cla() 
        image2 = image2[roi[0][0]:roi[1][0], roi[0][1]:roi[1][1]]
        sp2.imshow(image2,origin='lower')
        #print "Adding image2"
    
    def ShowFinalResults(self):
        """
        After registration has run, this function will display the starting pose
        final pose and metric results from optimizer.
        """
        # Get fixed image
        fixed_image = JT.utils.vtk2numpy(self._fixedImage)
        # Moving image will be found at initial pose and final pose
        initial_pose = self._initialStartingPose
        final_pose = self.bestPose
        # Ginerate initial image
        self._imageGenerator._volume.SetPose(initial_pose)
        self.UpdateImageGenerator()
        initial_moving = JT.utils.vtk2numpy(self._movingImage)
        # Generate final image
        self._imageGenerator._volume.SetPose(final_pose)
        self.UpdateImageGenerator()
        final_moving = JT.utils.vtk2numpy(self._movingImage)
        
        # Crop all images to ROI
        roi = self._transformedRoi
        fixed_image = fixed_image[roi[0][0]:roi[1][0], roi[0][1]:roi[1][1]]
        initial_moving = initial_moving[roi[0][0]:roi[1][0], roi[0][1]:roi[1][1]]
        final_moving = final_moving[roi[0][0]:roi[1][0], roi[0][1]:roi[1][1]]
        
        # Plot Results
        data = np.asarray(self._data)
        fig = pylab.figure()
        sp1a = pylab.subplot(431)
        sp1a.imshow(fixed_image,origin='lower')
        pylab.xticks([])
        pylab.yticks([])
        pylab.xlabel('Fixed Image')
        sp1b = pylab.subplot(432)
        sp1b.imshow(initial_moving,origin='lower')
        pylab.xticks([])
        pylab.yticks([])
        pylab.xlabel('Initial Image')
        sp1c = pylab.subplot(433)
        sp1c.imshow(final_moving,origin='lower')
        pylab.xticks([])
        pylab.yticks([])
        pylab.xlabel('Final Image')
        
        pylab.draw()
        sp2 = pylab.subplot(412)
        sp2.plot(data[:,2],'b-',
                   data[:,3],'g-',
                   data[:,4],'r-')
        pylab.ylabel('Translations (mm)')
        sp3 = pylab.subplot(413)
        sp3.plot(data[:,5],'b-',
                   data[:,6],'g-',
                   data[:,7],'r-')
        pylab.ylabel('Rotations (deg)')
        sp4 = pylab.subplot(414)
        sp4.plot(data[:,1])
        sp4.plot(data[:,8],'g-', lw=2)
        pylab.ylabel('Function Value')
        pylab.suptitle('Registration Results',fontsize=16)
        pylab.show()
        
    
    def Reset(self):
        """
        Reinitialize the registration class for a new run.
        """
        self.bestCost = 1e300
        self._iterCount = 0
        
    
    def Update(self):
        """
        Check to make sure that everything is up-to-date.
        Items to check:
            a fixed image has been set
            an moving image generator (drr) has been set
            roi should be scaled by the resolution factor
            drr and fixed image should have same pixel spacing and size
            
        """
        if self._fixedImage == None:
            print "Fixed image is not set"
            return False
        
        if not isinstance(self._fixedImage,JT.vtk.vtkImageData):
            print "Fixed image is not a VTK image"
            return False
        
        if self._movingImage == None:
            print "Moving image is not set"
            return False
        
        if not isinstance(self._movingImage,JT.vtk.vtkImageData):
            print "Moving image is not a VTK image"
            return False
        
        if self._fixedImage.GetDimensions() != self._movingImage.GetDimensions():
            print "Fixed image and moving image are not the same size."
            return False
        
        # Make sure that the moving image spacing matches the fixed image
        try:
            self._movingImage.SetSpacing(self._fixedImage.GetSpacing()) 
        except:
            print "Moving image spacing could not be set to match fixed image"
            return False                
        # Make sure that the ROI has been transformed to match the resolution factor
#        try:
#            self._transformedRoi = ((np.asarray(self._roi) /
#                                    self._resolutionFactor).astype('int'))
#            self._metricParams['roi'] = self._transformedRoi
#        except:
#            print "ROI could not be transformed"
#            return False
        
        # If we made it to here, everything must be ok return True
        return True
        


class MultiResolutionRegistration(Registration):
    """
    Multi resolution registration class.  
    """
    
    def __init__(self):
        """ Constructor
        """
        Registration.__init__(self)
        self._resolutionLevels = 4
        self._currentLevel = 4
        self._reductionFactor = 0.5
        self._optimizedPoses = []
        # EXPERIMENTAL: it seems like shrinking the images may not help???
        self._shrinkImages = False
        
    
    def SetResolutionLevels(self,levels):
        """
        Levels is an integer that represents the number of levels to 
        perform for the multi-resolution registration. Default is 4.
        
        The step size for the optimizer will be reduced by 
        self._reductionFactor at each level.  self._reductionFactor
        is set to a default of 0.5.  
        
        For example: for levels=4, reductionFactor=0.5, and startingStep = 4:
            level 1: step size = 4
            level 2: step size = 2
            level 3: step size = 1
            level 4: step size = 0.5 
        """
        self._resolutionLevels = levels
        self._currentLevel = levels
        
    def SetReductionFactor(self,factor):
        """
        Factor is a float in the range (0.0-1.0] that will be multiplied 
        to the current step size in the optimizer to determine the new 
        step size for the next level of registration. Default is 0.5.
        
        For example: for levels=4, reductionFactor=0.5, and startingStep = 4:
            level 1: step size = 4
            level 2: step size = 2
            level 3: step size = 1
            level 4: step size = 0.5 
        """
        if (factor <= 1.0) and (factor > 0.0):
            self._reductionFactor = factor
        else:
            print "Error. Factor must be in the range (0.0-1.0]"
            raise(self)      
            
    
    def GetValue(self,pose,volume=0,log_result=True):
        """
        return the current value of the Gradient Difference image metric.  Set the position and
        orientation of the transform.  
        INPUTS:
            Pose - vector made up of (position[3],orientation[3])
            volume - the (integer) index of item_list, to which the pose will be applied 
                     allows multiple volumes to be registered to the same fixed image.
            log_result - bool.  allows debug and logging output to be turned off.
        OUTPUT:
            value - scalar cost function evaluation.  Cost function can be set 
                    using my Metric.py set of functions or to any other function
                    that can take two images and return a scalar cost. 
        """
        global time_last

        #Check to make sure a 6 element pose was supplied
        try:
            assert len(pose)==6
        except AssertionError:
            print "Expected 'pose' to contain 6 elements, got %i" % len(pose)
            raise
    
        # First update the pose of the volume before the image is generated
        # WARNING: This interface might be changed to better manage multiple volumes
        self._imageGenerator._volume.SetPose(pose)
        self.UpdateImageGenerator()
        
        image1 = self._fixedImage
        image2 = self._movingImage
        
        # Check components to make sure everything is up-to-date
        # Returns false if there is a problem with the components
        if self.Update() == False:
            print "Could not update components, nothing done"
            raise AssertionError
        
        # Shrink the images for multiresolution 
        if self._shrinkImages:
            image1 = self.ShrinkImage(image1)
            image2 = self.ShrinkImage(image2)
        
        # Convert VTK images to numpy arrays for image metric calculation
        image1 = JT.utils.vtk2numpy(image1)
        image2 = JT.utils.vtk2numpy(image2)
        
        # This is the key metric evaluation step
        try:
            val=self._imageMetric.GetValue(image1,image2,**self._metricParams)
            self._iterCount += 1
        except:
            print "Error calculating metric"
            raise
        
        if val < self.bestCost:
            self.bestCost = val
            self.bestPose = pose

        if self._debug and log_result:
            print "%i(%fs), Metric:,%4.4f, Pose[mm,deg]:, %4.2f, %4.2f, %4.2f, %4.4f, %4.4f, %4.4f"%\
                   (self._iterCount,time.time()-time_last,val,pose[0],pose[1],
                    pose[2],np.degrees(pose[3]),np.degrees(pose[4]),np.degrees(pose[5]))
            time_last = time.time()
            # Log status/results to a file that can be plotted later.
            my_tupel = [self._iterCount, val, pose[0], pose[1], pose[2],
                        np.degrees(pose[3]), np.degrees(pose[4]),
                        np.degrees(pose[5]), self.bestCost]
            self.log.writerow(my_tupel)
        if self._plot:
            self.ShowImages(image1,image2)

        return val
    
    def Start(self):
        """
        Start the registration process by kicking off the optimizer.  For 
        Multi-resolution registration, the registration process will be 
        repeated in a coarse to fine manner, changing the appropriate
        optimizer parameters at each level.
        
        The optimizer parameter that will be adjusted are:
            MaximumStepLength: inversely proportional to level
            ***Might want to change scales***
            
        The additional Registration parameters that will be adjusted:
            ImageSize (shrink Image factor): inversely proportional to level
        """
        initialStepLength = self.optimizer.GetMaximumStepLength()
        self.Reset()
        for level in range(int(self._resolutionLevels),0,-1):
            print "Registration level: ", level
            self._currentLevel = level
            self.StartNewLog()
            # Do optimization with current settings
            self.optimizer.StartOptimization()
            self.OptimizerDone()
            self.Reset()
            # Take results from this level and setup next level
            self._optimizedPoses.append(self.bestPose)
            self.SetInitialPose(self.bestPose)
            current_step_length = self.optimizer.GetMaximumStepLength()
            step_length = current_step_length * self._reductionFactor
            self.optimizer.SetMaximumStepLength(step_length)
        
    
    def ShrinkImage(self,image,factor=None):
        """
        This function will subsample a 2D image so that it will be 
        smaller.  It is intended to be used in coarse to fine registration
        scheme.  Set factor to 1 to output original image size.  
        
        image is a VTK image
        factor [optional] - shrink factor.  
                            Will be the same in both X and Y dimensions.
        if no factor is given, the image will be shrunk using inverse of 
        the current resolution level as the shrink factor.  For example if 
        the current resolution level is 4, the image will be shrunk to 1/4 size.
        
        returns the shrunk image in VTK format.
        """
        if factor == None:
            try:
                factor = int(self._currentLevel)
            except:
                print "Could not shrink image using resolution level: ",\
                      self._currentLevel
                print "Using full size image instead."
                factor = 1
                
        shrinker = JT.vtk.vtkImageShrink3D()    
        shrinker.SetShrinkFactors(factor,factor,1)
        shrinker.SetInput(image)
        shrinker.Update()
        return shrinker.GetOutput()
    
    def SetRegionOfInterest(self,roi):
        """
        Set a region of interest for to apply the metric to.  This could be 
        the same region specified in a image mask.  This roi will be transformed
        (scaled) by the resolution factor so that it will match the image at
        the desired resolution.
        """
        self._roi = roi
        # roi needs to be scaled to fit the image at different resolution factors
        self._transformedRoi = np.asarray(roi) / self._currentLevel
        self._metricParams['roi'] = self._transformedRoi
    
    def Update(self):
        """
        Check to make sure that everything is up-to-date.
        Items to check:
            a fixed image has been set
            an moving image generator (drr) has been set
            roi should be scaled by the resolution factor
            drr and fixed image should have same pixel spacing and size
            
        """
        if self._fixedImage == None:
            print "Fixed image is not set"
            return False
        
        if not isinstance(self._fixedImage,JT.vtk.vtkImageData):
            print "Fixed image is not a VTK image"
            return False
        
        if self._movingImage == None:
            print "Moving image is not set"
            return False
        
        if not isinstance(self._movingImage,JT.vtk.vtkImageData):
            print "Moving image is not a VTK image"
            return False
        
        if self._fixedImage.GetDimensions() != self._movingImage.GetDimensions():
            print "Fixed image and moving image are not the same size."
            return False
        
        # Make sure that the moving image spacing matches the fixed image
        try:
            self._movingImage.SetSpacing(self._fixedImage.GetSpacing()) 
        except:
            print "Moving image spacing could not be set to match fixed image"
            return False                
        
        # Make sure that the ROI has been transformed to match the resolution factor
        if self._shrinkImages:
            try:
                self._transformedRoi = ((np.asarray(self._roi) /
                                        self._currentLevel).astype('int'))
                self._metricParams['roi'] = self._transformedRoi
            except:
                print "ROI could not be transformed"
                return False
        else:
            self._transformedRoi = np.asarray(self._roi)
            self._metricParams['roi'] = self._transformedRoi
                    
        # If we made it to here, everything must be ok return True
        return True

class MultiIterationRegistration(Registration):
    """
    A MultiIteration Registration Class.  Specify the number of times the registration
    should run, using the previous result as the starting point for the next 
    registration Iteration.  Unlike the multiresolution registration, all 
    registration is done on full scale images.
    """
    
    def __init__(self):
        """
        """
        Registration.__init__(self)
        self._iterations = 2
        self._initialStepLength = 4.0
        self._reductionFactor = 0.6180339887
        self._currentIteration = 1
        self._optimizedPoses = []
        
        
    def SetIterations(self,iterations):
        """
        Iterations is an integer that specifies how many times the registration 
        should run.  Setting Iterations to 1 will have the same effect as 
        using the Registration base class.
        """
        self._iterations = iterations
        
    def SetInitialStepLength(self,step_length):
        """
        float. The step length that should be used initially, this will be reduced 
        with each iteration.
        """
        self._initialStepLength = step_length
        self.optimizer.SetMaximumStepLength(step_length)
        
    
    def SetReductionFactor(self,factor):
        """
        Factor is a float in the range (0.0-1.0] that will be multiplied 
        to the current step size in the optimizer to determine the new 
        step size for the next iteration of registration. Default is 0.61.
        
        For example: for Iterations=4, reductionFactor=0.5, and startingStep = 4:
            Iteration 1: step size = 4
            Iteration 2: step size = 2
            Iteration 3: step size = 1
            Iteration 4: step size = 0.5 
        """
        if (factor <= 1.0) and (factor > 0.0):
            self._reductionFactor = factor
        else:
            print "Error. Factor must be in the range (0.0-1.0]"
            raise(self) 
    
    def Start(self):
        """
        Start the registration process by kicking off the optimizer.  For 
        Multi-Iteration registration, the registration process will be 
        repeated for the number of times specified by Iterations.
        
        The optimizer parameter that will be adjusted are:
            Step length - decreased by Reduction Factor for each reg Iteration
        """
        # Append to the same log for all iterations
        self.StartNewLog()
        self.Reset()
        self.optimizer.SetMaximumStepLength(self._initialStepLength)
        for iter in range(int(self._iterations)):
            print "Registration Iteration: ", iter+1
            print "Step Length: ", self.optimizer.GetMaximumStepLength()
            self._currentIteration = iter
            # Do optimization with current settings
            self.optimizer.StartOptimization()
            self.IterationDone()
            # Take results from this level and setup next level
            self._optimizedPoses.append(self.bestPose)
            self.SetInitialPose(self.bestPose)
            current_step_length = self.optimizer.GetMaximumStepLength()
            step_length = current_step_length * self._reductionFactor
            self.optimizer.SetMaximumStepLength(step_length)
        
        print "Registration Finished!"
        self.OptimizerDone()
        
    def IterationDone(self):
            """
            Perform any clean up functions necessary after the optimizer is 
            finished running, but don't close the log file (as in OptimizerDone)
            """
            # Make sure that the registration did not find a better solution 
            # while doing the gradient descent calculations
            if self.bestCost < self.optimizer._Value:
                print "The optimizer returned a value of: ", self.optimizer._Value
                print "But this was not the best solution found."
                print "Which was: ", self.bestCost
                print "At Pose  : ", self.bestPose
            else:
                self.bestCost = self.optimizer._Value
                self.bestPose = self.optimizer._CurrentPosition


class BiPlaneRegistration(Registration):
    """
    This registration class enables biplane registration (2 fluoro images
    and one Volume).  For this to work, 2 DRR's need to be setup, one to 
    generate moving images for each of the fluoro images.
    """
    
    def __init__(self):
        """
        """
        Registration.__init__(self)
        self._imageGenerator = []
        self._fixedImage = []
        self._movingImage = []
        self._originalFixedImageSize = []
        
    
    def Start(self):
        pass
        
    
    def AddMovingImageGenerator(self,ImageGenerator):
        """
        Add an object that can generate a 2D moving image.  This object
        will be appended to an imageGenerator list that must correspond to 
        the fixed image list so that the correct fixed and moving images
        can be registered.
        Currently, this is limited to the DRR class.
        IMPORTANT: fixed image must be set first so that the DRR image size
        can be properly set.  
        """
        # Since the fixed image should be set first, there should be
        # one more item in the fixed image list than the moving image list 
        if len(self._fixedImage) != (len(self._imageGenerator) - 1):
            print "Error: Fixed image must be set first"
            return
        
        # Add the supplied image generator to the list
        self._imageGenerator.apppend(ImageGenerator)
        self._imageGenerator[-1].SetRenderWindowSize(self._originalFixedImageSize[-1])
        
        # nPlanes is the number of imageGenerators, each of which can represent 
        # a different projection plane of the volume.
        self.nProjections = len(self._imageGenerator)
        
        # make moving image list the same length as the image generator list
        self._movingImage = range(self.nProjections)
        
        # Return the current number of imageGenerators
        return self.nProjections


    def UpdateImageGenerator(self):
        """
        Call the GenerateImage function and save the resulting image as the 
        new moving image, for all image generators in the imageGenerator list.
        """
        for i,generator in enumerate(self._imageGenerator):
            self._movingImage[i] = generator.GetImage()
            
            #Set the moving image spacing to be the same as fixed image spacing
            #Is this really necessary each time an image is generated???
            self._movingImage[i].SetSpacing(self._fixedImage[i].GetSpacing())

    
    def AddFixedImage(self,fixedImage):
        """
        Set a VTK image.  The expected usage is:
        reg.AddFixedImage(FixedImage.GetImage(i))
        
        It will be converted to a numpy image in the GetValue method.
        """
        self._fixedImage.append(fixedImage)
        self._originalFixedImageSize.append(fixedImage.GetDimensions()[0:2])
        
        self.nFixedImages = len(self._fixedImage)
        return self.nFixedImages
        
        
    def GetValue(self,pose,volume=0,log_result=True):
        """
        return the current value of the image metric.  Set the position and
        orientation of the transform.  
        INPUTS:
            Pose - vector made up of (position[3],orientation[3])
            volume - the (integer) index of item_list, to which the pose will be applied 
                     allows multiple volumes to be registered to the same fixed image.
        OUTPUT:
            value - scalar cost function evaluation.  Cost function can be set 
                    using my Metric.py set of functions or to any other function
                    that can take two images and return a scalar cost. 
        """
        global time_last

        #Check to make sure a 6 element pose was supplied
        try:
            assert len(pose)==6
        except AssertionError:
            print "Expected 'pose' to contain 6 elements, got %i" % len(pose)
            raise
    
        # First update the pose of the volume before the image is generated
        # WARNING: This interface might be changed to better manage multiple volumes
        self._imageGenerator._volume.SetPose(pose)
        self.UpdateImageGenerator()
        
        image1 = self._fixedImage
        image2 = self._movingImage
        
        # Check components to make sure everything is up-to-date
        # Returns false if there is a problem with the components
        if self.Update() == False:
            print "Could not update components, nothing done"
            raise AssertionError
        
        # Convert VTK images to numpy arrays for image metric calculation
        image1 = JT.utils.vtk2numpy(image1)
        image2 = JT.utils.vtk2numpy(image2)
        
        # This is the key metric evaluation step
        try:
            val=self._imageMetric.GetValue(image1,image2,**self._metricParams)
            self._iterCount += 1
        except:
            print "Error calculating metric"
            raise
                
        if val < self.bestCost:
            self.bestCost = val
            self.bestPose = pose

        if self._debug and log_result:
            # only write results if debug is set and log_result is true
            # log_result is turned off when gradients are calculated
            print "%i(%fs), Metric:,%4.4f, Pose[mm,deg]:, %4.2f, %4.2f, %4.2f, %4.4f, %4.4f, %4.4f"%\
                   (self._iterCount,time.time()-time_last,val,pose[0],pose[1],
                    pose[2],np.degrees(pose[3]),np.degrees(pose[4]),np.degrees(pose[5]))
            time_last = time.time()
            # Log status/results to a file that can be plotted later.
            data_list = [self._iterCount, val, pose[0], pose[1], pose[2],
                        np.degrees(pose[3]), np.degrees(pose[4]),
                        np.degrees(pose[5]), self.bestCost]
            self._data.append(data_list)
            self.log.writerow(data_list)
        
        if self._plot and log_result:
            self.ShowImages(image1,image2)
        

        return val

        
                  